#!/usr/bin/env python
# -*- coding:utf-8 -*-
from unicorn import *
from unicorn.x86_const import *

import os,sys,struct
from lib.utils import printer

class Register:
    registers = {
        "X86":{
            'ax': UC_X86_REG_EAX,
            'bx': UC_X86_REG_EBX,
            'cx': UC_X86_REG_ECX,
            'dx': UC_X86_REG_EDX,
            'di': UC_X86_REG_EDI,
            'si': UC_X86_REG_ESI,
            'bp': UC_X86_REG_EBP,
            'sp': UC_X86_REG_ESP,
            'ip': UC_X86_REG_EIP,
        },
        "X64":{
            'ax': UC_X86_REG_RAX,
            'bx': UC_X86_REG_RBX,
            'cx': UC_X86_REG_RCX,
            'dx': UC_X86_REG_RDX,
            'di': UC_X86_REG_RDI,
            'si': UC_X86_REG_RSI,
            'bp': UC_X86_REG_RBP,
            'sp': UC_X86_REG_RSP,
            'ip': UC_X86_REG_RIP,
        },
        "eflags":{
            'CF[进位标志]': '0',
            'PF[奇偶标志]': '2',
            'AF[辅助进位]': '4',
            'ZF[零标志位]': '6',
            'SF[符号标志]': '7',
            'TF[跟踪标志]': '8',
            'IF[中断标志]': '9',
            'DF[方向标志]': '10',
            'OF[溢出标志]': '11'
        }
    }

    def __init__(self,mu,arch):
        self.mu = mu
        self.arch = arch
    
    def print_registers(self):
        if self.arch == 'X86':
            printer.test('eax: %.8X ebx: %.8X ecx: %.8X edx: %.8X' % (
                                self.mu.reg_read(UC_X86_REG_EAX),
                                self.mu.reg_read(UC_X86_REG_EBX), 
                                self.mu.reg_read(UC_X86_REG_ECX), 
                                self.mu.reg_read(UC_X86_REG_EDX)
                            ),flag=''
                        )
            printer.test('esp: %.8X ebp: %.8X esi: %.8X edi: %.8X' % (
                                self.mu.reg_read(UC_X86_REG_ESP), 
                                self.mu.reg_read(UC_X86_REG_EBP), 
                                self.mu.reg_read(UC_X86_REG_ESI), 
                                self.mu.reg_read(UC_X86_REG_EDI)
                            ),flag=''
                        )
                            
            printer.test('eip: %.8X' % (
                                self.mu.reg_read(UC_X86_REG_EIP)
                            ),flag=''
                        )
        elif self.arch == "X64":
            printer.test('rax: %.8X ebx: %.8X ecx: %.8X edx: %.8X' % (
                                self.mu.reg_read(UC_X86_REG_RAX), 
                                self.mu.reg_read(UC_X86_REG_RBX), 
                                self.mu.reg_read(UC_X86_REG_RCX), 
                                self.mu.reg_read(UC_X86_REG_RDX)
                            ),flag=''
                        )
                            
            printer.test('rsp: %.8X rbp: %.8X rsi: %.8X rdi: %.8X' % (
                                self.mu.reg_read(UC_X86_REG_RSP), 
                                self.mu.reg_read(UC_X86_REG_RBP), 
                                self.mu.reg_read(UC_X86_REG_RSI), 
                                self.mu.reg_read(UC_X86_REG_RDI)
                            ),flag=''
                        )
                            
            printer.test('rip: %.8X' % (
                                self.mu.reg_read(UC_X86_REG_RIP)
                            ),flag=''
                        )
        printer.test('fs: %.8X gs: %.8X cs: %.8X  ds: %.8X  es: %.8X' % (
                            self.mu.reg_read(UC_X86_REG_FS), 
                            self.mu.reg_read(UC_X86_REG_GS), 
                            self.mu.reg_read(UC_X86_REG_CS), 
                            self.mu.reg_read(UC_X86_REG_DS), 
                            self.mu.reg_read(UC_X86_REG_ES)
                        ),flag=''
                    )
        printer.test("-------------------- Eflags[标志位寄存器信息] --------------------",flag='')
        self.print_eflags()
        printer.test("-----------------------------------------------------------------",flag='')
    def print_eflags(self):
        eflags = '{:016b}'.format(self.mu.reg_read(UC_X86_REG_EFLAGS))
        eflags = eflags[::-1]
        for i in range(len(eflags)):
            if self.registers['eflags']['CF[进位标志]'] == str(i):
                #registers['eflags']['CF[进位标志]'] = eflags[i]
                print(" CF[进位标志] = " + eflags[i],end='  ')
            elif self.registers['eflags']['PF[奇偶标志]'] == str(i):
                #registers['eflags']['PF[奇偶标志]'] = eflags[i]
                print(" PF[奇偶标志] = " + eflags[i],end='  ')
            elif self.registers['eflags']['AF[辅助进位]'] == str(i):
                #registers['eflags']['AF[辅助进位标志]'] = eflags[i]
                print(" AF[辅助进位] = " + eflags[i],end='  ')
            elif self.registers['eflags']['ZF[零标志位]'] == str(i):
                #registers['eflags']['ZF[零标志]'] = eflags[i]
                print(" ZF[零标志位] = " + eflags[i])
            elif self.registers['eflags']['SF[符号标志]'] == str(i):
                #registers['eflags']['SF[符号标志]'] = eflags[i]
                print(" SF[符号标志] = " + eflags[i],end='  ')
            elif self.registers['eflags']['TF[跟踪标志]'] == str(i):
                #registers['eflags']['TF[跟踪标志]'] = eflags[i]
                print(" TF[跟踪标志] = " + eflags[i],end='  ')
            elif self.registers['eflags']['IF[中断标志]'] == str(i):
                #registers['eflags']['IF[中断标志]'] = eflags[i]
                print(" IF[中断标志] = " + eflags[i],end='  ')
            elif self.registers['eflags']['DF[方向标志]'] == str(i):
                #registers['eflags']['DF[方向标志]'] = eflags[i]
                print(" DF[方向标志] = " + eflags[i])
            elif self.registers['eflags']['OF[溢出标志]'] == str(i):
                #registers['eflags']['OF[溢出标志]'] = eflags[i]
                print(" OF[溢出标志] = " + eflags[i])


